import cv2
import argparse
import os
import math
import sys
from pathlib import Path
import numpy as np
 
sys.path.append(str(Path(__file__).resolve().parent.parent))

from file_util.file_util import GetAllSubFile
import xml.etree.ElementTree as ET

def update_anno(anno_file, w, h, new_anno_file, angle):
		try:
			tree = ET.parse(anno_file)
			root = tree.getroot() 
		except:
			print(anno_file + " is not exist.")
			return
		angle = angle / 180 * np.pi
		for object in root:
			if object.tag == "object":
				bndbox = object[4][0]
				xmin = int(object[4][0].text)
				ymin = int(object[4][1].text)
				xmax= int(object[4][2].text)
				ymax = int(object[4][3].text)
				print("xmin=",xmin,"ymin=",ymin,"xmax=",xmax,"ymax=",ymax)
				centerx= (xmax - xmin)/2 + xmin
				centery= (ymax - ymin)/2 + ymin
				print("centerx=",centerx,"centery=",centery)
				center_xmin = -(xmax - xmin)/2
				center_ymin = (ymax - ymin)/2
				center_xmax = (xmax - xmin)/2
				center_ymax = -(ymax - ymin)/2
				print("center_xmin=",center_xmin,"center_ymin=",center_ymin,"center_xmax=",center_xmax,"center_ymax=",center_ymax)
				center_new_xmin = center_xmin * np.cos(angle) - center_ymin * np.sin(angle)
				center_new_ymin = center_ymin * np.cos(angle) + center_xmax * np.sin(angle)
				center_new_xmax = center_xmin * np.cos(angle) - center_ymax * np.sin(angle)
				center_new_ymax = center_ymax * np.cos(angle) + center_xmax * np.sin(angle)
				print("center_new_xmin=",center_new_xmin,"center_new_ymin=",center_new_ymin,"center_new_xmax=",center_new_xmax,"center_new_ymax=",center_new_ymax)
				new_xmin = int(centerx - center_new_xmin)
				new_ymin = int(centery + center_new_ymin)
				new_xmax = int(centerx - center_new_xmax)
				new_ymax = int(centery + center_new_ymax)
				print("new_xmin=",new_xmin,"new_ymin=",new_ymin,"new_xmax=",new_xmax,"new_ymax=",new_ymax)
				object[4][0].text = str(new_xmin)
				object[4][1].text = str(new_ymin)
				object[4][2].text = str(new_xmax)
				object[4][3].text = str(new_ymax)
		tree.write(new_anno_file,encoding="utf-8", xml_declaration=True)
		return 
def update_anno1(anno_file, centerx, centery,w,h, new_anno_file, angle,org_angle):
		try:
			tree = ET.parse(anno_file)
			root = tree.getroot() 
		except:
			print(anno_file + " is not exist.")
			return
		remove_objects = []
		for object in root:
			if object.tag == "object":
				bndbox = object[4][0]
				xmin = int(object[4][0].text)
				ymin = int(object[4][1].text)
				xmax= int(object[4][2].text)
				ymax = int(object[4][3].text)
				#centerx= (xmax + xmin)/2 
				#centery= (ymax + ymin)/2 
				new_x_topleft, new_y_topleft = getPointAffinedPos(xmin,ymin,centerx,centery,angle)
				new_x_bottomright, new_y_bottomright = getPointAffinedPos(xmax,ymax,centerx,centery,angle)
				new_x_topright, new_y_topright = getPointAffinedPos(xmax,ymin,centerx,centery,angle)
				new_x_bottomleft, new_y_bottomleft = getPointAffinedPos(xmin,ymax,centerx,centery,angle)
				
				if (org_angle > 0) and (org_angle <= 90):
					object[4][0].text = str(new_x_topleft)
					object[4][1].text = str(new_y_topright)
					object[4][2].text = str(new_x_bottomright)
					object[4][3].text = str(new_y_bottomleft)
				elif (org_angle > 90) and (org_angle <= 180):
					object[4][0].text = str(new_x_topright)
					object[4][1].text = str(new_y_bottomright)
					object[4][2].text = str(new_x_bottomleft)
					object[4][3].text = str(new_y_topleft)
				if (org_angle > 180) and (org_angle <= 270):
					object[4][0].text = str(new_x_bottomright)
					object[4][1].text = str(new_y_bottomleft)
					object[4][2].text = str(new_x_topleft)
					object[4][3].text = str(new_y_topright)
				elif (org_angle > 270) and (org_angle <= 360):
					object[4][0].text = str(new_x_bottomleft)
					object[4][1].text = str(new_y_topleft)
					object[4][2].text = str(new_x_topright)
					object[4][3].text = str(new_y_bottomright)
				new_xmin = int(object[4][0].text)
				new_ymin = int(object[4][1].text)
				new_xmax= int(object[4][2].text)
				new_ymax = int(object[4][3].text)
				if (new_xmax < 0) or (new_ymax < 0):
					print(object[0].text)
					#root.remove(object)
					remove_objects.append(object)
					continue
				if (new_xmin > w) or (new_ymin > h):
					print(object[0].text)
					#root.remove(object)
					remove_objects.append(object)
					continue
				if new_xmin < 0:
					new_xmin = 0
				if new_ymin < 0 :
					new_ymin = 0
				if new_xmax > w :
					new_xmax = w
				if new_ymax > h :
					new_ymax = h
				object[4][0].text = str(new_xmin)
				object[4][1].text = str(new_ymin)
				object[4][2].text = str(new_xmax)
				object[4][3].text = str(new_ymax)
		for remove_object in remove_objects:
			root.remove(remove_object)
		tree.write(new_anno_file,encoding="utf-8", xml_declaration=True)
		return 
def getPointAffinedPos(src_x, src_y, center_x, center_y, angle):
    x = src_x - center_x
    y = src_y - center_y
    dst_x = round(x * np.cos(angle) + y * np.sin(angle) + center_x)
    dst_y = round(-x * np.sin(angle) + y * np.cos(angle) + center_y)
    return dst_x, dst_y;

def draw_totation(new_img_file,new_anno_file):
		try:
			tree = ET.parse(new_anno_file)
			root = tree.getroot() 
		except:
			print(new_anno_file + " is not exist.")
			return
		img = cv2.imread(new_img_file)
		for object in root:
			if object.tag == "object":
				bndbox = object[4][0]
				xmin = int(object[4][0].text)
				ymin = int(object[4][1].text)
				xmax= int(object[4][2].text)
				ymax = int(object[4][3].text)
				cv2.rectangle(img,(xmin,ymin),(xmax,ymax),(0,0,255),5)
		cv2.imwrite(new_img_file,img)

def rotation_crop(rotation_img,w,h,centerx, centery, angle):
	new_x_topleft, new_y_topleft = getPointAffinedPos(0,0,centerx,centery,angle)
	new_x_bottomright, new_y_bottomright = getPointAffinedPos(w,h,centerx,centery,angle)
	new_x_topright, new_y_topright = getPointAffinedPos(w,0,centerx,centery,angle)
	new_x_bottomleft, new_y_bottomleft = getPointAffinedPos(o,h,centerx,centery,angle)

def data_rotation(opt,rotation_angle_list):
	img_files = GetAllSubFile(opt.root_data)
	for img_file in img_files:
		img = cv2.imread(img_file)
		h = img.shape[0]
		w = img.shape[1]
		center = (w/2, h/2)
		path, img_file_name = os.path.split(img_file)
		img_name = img_file_name.rsplit('.',1)[0]
		#draw_totation(img_file,os.path.join(opt.anon_dir,img_name + '.xml'))
		anno_file = os.path.join(opt.anon_dir,img_name + '.xml')
		for angle in rotation_angle_list:
			M = cv2.getRotationMatrix2D(center,angle,1.0)
			angle1 = angle / 180 * np.pi  
			#h1 = abs(round(w * np.sin(angle) + h * np.cos(angle)))
			#w1 = abs(round(w * np.cos(angle) + h * np.sin(angle)))
			#print("new_ w: ", w1, ' h:', h1)
			#M[0, 2] += (w1 - w) / 2
			#M[1, 2] += (h1 - h) / 2
			rotation_img = cv2.warpAffine(img, M, (w, h), borderValue=(114,114,114))
			#rotation_img = cv2.warpAffine(img, M, (w,h))
			new_img_name = img_name+"_rotation_" + str(angle)
			new_anno_file = os.path.join(opt.anon_dir,new_img_name + '.xml')
			new_img_file = os.path.join(path,new_img_name +".jpg")
			cv2.imwrite(new_img_file,rotation_img)
			print(new_anno_file)
			update_anno1(anno_file, center[0], center[1],w,h, new_anno_file, angle1, angle)
			#draw_totation(new_img_file,new_anno_file)
		
		

if __name__ == '__main__':
   parser = argparse.ArgumentParser()
   parser.add_argument('--root_data', type=str, default='./data', help='root data dir.')
   parser.add_argument('--anon_dir', type=str, default='./anno', help='anno dir.')
   opt = parser.parse_args()
   opt.anon_dir = opt.anon_dir.strip()
   rotation_angle_list = [30,60,90,120,150,180,210,240,270,300,330]
   #rotation_angle_list = [15,30,45,60,75,105,120,135,150,165]
   #rotation_angle_list = [15,30]
   #11
   data_rotation(opt,rotation_angle_list)